#本脚本的目的是在从origin-path读取一张图片，通过可视化界面、鼠标画笔来生成mask，并将它存在target-path

import argparse
from pathlib import Path
import cv2
import numpy as np
import os

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument('--origin-path', default=r"D:\tempDownload\dataset\elevator\glass\img\000_0120.jpg", type=str)
    parser.add_argument('--target-path', default=r"D:\data\repository\small-code", type=str)
    return parser.parse_args()

def Draw_Mask(event, x, y, flags, param=None):
    if event == cv2.EVENT_LBUTTONDOWN:
        param['lbutton_state'] = 1
        param['init_point'] = (x,y)
        param['previous_image_states'].append(param['image_mask'].copy())
    elif param['lbutton_state'] == 1 and flags == cv2.EVENT_FLAG_LBUTTON:
        cv2.line(param['image_mask'], param['init_point'], (x, y), (0, 0, 0), param["draw_size"])
        param['init_point'] = (x,y)
    elif event == cv2.EVENT_LBUTTONUP:
        param['lbutton_state'] = 0
        cv2.line(param['image_mask'], param['init_point'], (x, y), (0, 0, 0), param["draw_size"])
        cv2.imwrite(f"{param['target_img_dir']}{os.sep}{param['image_path'].stem}_mask.jpg", param['image_mask'])
        print(f"saveMaskSuccess!!! 放在{param['target_img_dir']}{os.sep}{param['image_path'].stem}_mask.jpg")
    elif event == cv2.EVENT_RBUTTONDOWN:
        if param['previous_image_states']:
            param['image_mask'] = param['previous_image_states'].pop()
        cv2.imwrite(f"{param['target_img_dir']}{os.sep}{param['image_path'].stem}_mask.jpg", param['image_mask'])
        print(f"saveMaskSuccess!!! 放在{param['target_img_dir']}{os.sep}{param['image_path'].stem}_mask.jpg")

    cv2.imshow(f"{param['image_path'].name}", cv2.bitwise_and(param['image'], param['image_mask']))  # 显示的是图片矩阵

def Get_Draw_Size(x, params):
    params['draw_size'] = x

def main():
    args = Parse_Arguments()
    origin_path = Path(args.origin_path.replace("\\", "/"))
    target_path = Path(args.target_path.replace("\\", "/"))
    image = cv2.imread(str(origin_path))
    origin_path_mask = Path(str(origin_path.parent) + '/' + origin_path.stem + "_mask.jpg")
    print(origin_path_mask)
    if origin_path_mask.exists():
        mask = cv2.imread(str(origin_path_mask))
        print(f"成功读入{origin_path_mask}mask图")
    else:
        mask = np.full((image.shape[0], image.shape[1], 3), 255, np.uint8)
        print(f"创建空白画布图")
    params = {
        'draw-size':20, #画笔尺寸
        'init_point':(0,0), #初始鼠标位置点  矩形拉框  画笔涂抹都需要
        'lbutton_state': 0, #左键是否被按下
        'previous_image_states': [],    #存储矩阵序列，可以撤回
        'image_mask': mask,      #存储中间矩阵，其实可以优化掉
        'image_rem': image.copy(),      #存储备份矩阵，画图都是使用该矩阵
        'image': image,                 #原图矩阵
        'image_path': origin_path,       #原图文件位置
        'target_img_dir':target_path #保存文件的目标位置
    }
    cv2.namedWindow(f'{origin_path.name}', cv2.WINDOW_NORMAL)
    cv2.moveWindow(f'{origin_path.name}', 600, 250)
    cv2.resizeWindow(f'{origin_path.name}', 1500, 980)
    cv2.createTrackbar('draw_size', f'{origin_path.name}', 10, 255, lambda x:Get_Draw_Size(x, params))
    cv2.setMouseCallback(f'{origin_path.name}', Draw_Mask, params)
    cv2.imshow(f'{origin_path.name}',image)
    cv2.waitKeyEx(0)
if __name__ == "__main__":
    main()
